#!/usr/local/bin/python3.7
import pyaudio
import wave
import os
import sys
import threading
from audiosplit import splitwav
from asr import asrwavs
from multiprocessing import Pipe

CHUNK = 512
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16*1024

# 默认180秒保存一次数据，并开始识别
RECORD_SECONDS = 180

# record_stat :0-continue,1-save and continue,2-save and end
record_stat = 0

split_tmp_ext = 1

psend,precv = Pipe()

output_texts = []

class myAsr (threading.Thread):
	def __init__(self):
		threading.Thread.__init__(self)
	def run(self):
		global split_tmp_ext
		global precv
		global output_texts
		while True:
			try:
				filename = precv.recv()
			except:
				print('end.')
				break
			print(f'开始识别：{filename}')
			names = splitwav(filename.strip(), split_tmp_ext)
			split_tmp_ext += 1
			if len(names)==0:
				print('wave file split fail.')
				os.exit(0)
			res = asrwavs(names)
			print( res )
			output_texts.append(res)
			#清理临时文件和文件夹
			for i in range(len(names)):
				os.unlink(names[i])
			os.rmdir( os.path.dirname(names[0]) )

class myThread (threading.Thread):
	def __init__(self):
		threading.Thread.__init__(self)
	def run(self):
		threads = recordaudio()
		for i in range(len(threads)):
			threads[i].join()

def control_by_key():
	global record_stat
	while True:
		print('按回车键保存当前音频;')
		print('输入\'q\'再按回车键保存并结束程序。')
		s = sys.stdin.readline()
		if s.strip() == '':
			record_stat = 1
			print('你按了回车键')
		elif s.strip() == 'q':
			record_stat = 2
			print('你输入了\'q\'，退出程序')
			break
			
def savewav( frames:list, filename:str ):
	wf = wave.open(filename, 'wb')
	wf.setnchannels(CHANNELS)
	wf.setsampwidth(pyaudio.get_sample_size(FORMAT))
	wf.setframerate(RATE)
	wf.writeframes(b''.join(frames))
	wf.close()

def recordaudio() -> list:
	global record_stat
	global port
	
	p = pyaudio.PyAudio()
	stream = p.open(format=FORMAT,
				channels=CHANNELS,
				rate=RATE,
				input=True,
				frames_per_buffer=CHUNK)

	print("recording...")

	frames = []
	idx = 1
	threads = []
	global psend
	while record_stat < 2:
		for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
			if record_stat > 0:
				break
			data = stream.read(CHUNK)
			frames.append(data)
		if record_stat == 1:
			record_stat = 0
		savewav(frames, f'rec/a{idx:04}.wav')
		print(f'save wav rec/a{idx:04}.wav')
		psend.send(f'rec/a{idx:04}.wav\n')
		idx+=1
		frames.clear()
	psend.close()
	
	# Stop and close the stream 
	stream.stop_stream()
	stream.close()
	# Terminate the PortAudio interface
	p.terminate()
	return threads

if __name__=='__main__':
	try:
		os.mkdir('rec')
	except:
		None
	t1 = myThread()
	t2 = myAsr()
	t1.start()
	t2.start()
	control_by_key()
	
	t1.join()
	t2.join()
	output_name = 'output.txt'
	if len(sys.argv)==2:
		output_name = sys.argv[1]
	fp = open( output_name, 'w' )
	for i in range(len(output_texts)):
		fp.write(output_texts[i])
	fp.close()
